# 识别物体颜色
import cv2
import numpy as np

# ------------------------------参数集合------------------------------
trackbar_default_minh = 84  # 滚动条默认值h 156 180  84 107
trackbar_default_maxh = 107
trackbar_default_mins = 43  # 滚动条默认值s
trackbar_default_maxs = 255
trackbar_default_minv = 46  # 滚动条默认值v
trackbar_default_maxv = 255
ball_color = 'pink'
# minp = np.array([156, 43, 46])  # 粉色HSV范围
# # maxp = np.array([180, 255, 255])
# 控制点在世界坐标系的坐标
objPoints = np.array([[0, 0, 0],
                      [0, 20, 0],
                      [20, 20, 0],
                      [-20, 0, 0]], dtype=np.float64)
# 相机内参矩阵与外参矩阵
cameraMatrix = np.array([[724.7037, 0, 356.4751],
                         [0, 721.6658, 319.4040],
                         [0, 0, 1]], dtype=np.float64)
distCoeffs = np.array([-0.4244, 0.2585, 0, 0, 0], dtype=np.float64)

# ------------------------------创建窗口------------------------------
windowname = 'frame'
cv2.namedWindow(windowname, cv2.WINDOW_AUTOSIZE)


# ------------------------------滚动条回调函数------------------------------
def P():
    pass


cv2.createTrackbar('minh', windowname, trackbar_default_minh, 255, P)
cv2.createTrackbar('maxh', windowname, trackbar_default_maxh, 255, P)
cv2.createTrackbar('mins', windowname, trackbar_default_mins, 255, P)
cv2.createTrackbar('maxs', windowname, trackbar_default_maxs, 255, P)
cv2.createTrackbar('minv', windowname, trackbar_default_minv, 255, P)
cv2.createTrackbar('maxv', windowname, trackbar_default_maxv, 255, P)


# ------------------------------识别颜色函数------------------------------
def color(frame):
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)     # 转化成HSV图像

    minh = cv2.getTrackbarPos('minh', windowname)
    maxh = cv2.getTrackbarPos('maxh', windowname)
    mins = cv2.getTrackbarPos('mins', windowname)
    maxs = cv2.getTrackbarPos('maxs', windowname)
    minv = cv2.getTrackbarPos('minv', windowname)
    maxv = cv2.getTrackbarPos('maxv', windowname)
    ranges = cv2.inRange(hsv, np.array([minh, mins, minv]), np.array([maxh, maxs, maxv]))
    gs_frame = cv2.GaussianBlur(ranges, (5, 5), 0)      # 高斯模糊
    erode_hsv = cv2.erode(gs_frame, None, iterations=2)    # 先腐蚀再膨胀
    dilate = cv2.dilate(erode_hsv.copy(), None, iterations=2)

    cnts = cv2.findContours(dilate.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
    approx = []
    if cnts:
        approx, frame = point(cnts, frame)
    cv2.imshow('dilate', dilate)
    cv2.imshow('img', frame)
    return approx


# ------------------------------拟合轮廓寻找角点函数------------------------------
def point(cnts, frame):
    approx = []
    c = max(cnts, key=cv2.contourArea)
    ((x, y), radius) = cv2.minEnclosingCircle(c)  # 找到最小的圆包含这个轮廓，返回坐标和半径
    if radius > 20:  # 判断轮廓的半径，太小的话就认为是噪声，就忽略掉
        approx = cv2.approxPolyDP(c, 40, True)  # 拟合轮廓得到角点
        if len(approx) == 4:
            for i in range(len(approx)):
                frame = cv2.line(frame, (approx[i][0][0], approx[i][0][1]),
                                 (approx[(i + 1) % 4][0][0], approx[(i + 1) % 4][0][1]), (255, 0, 0), 3)
                cv2.circle(frame, (approx[i][0][0], approx[i][0][1]), 4, (0, 0, 255), -1)
    return approx, frame


# ------------------------------pnp算法测距函数------------------------------
def pnp(approx):
    # 检测到的角点的坐标
    imgPoints = np.array([[approx[0][0][0], approx[0][0][1]],
                          [approx[1][0][0], approx[1][0][1]],
                          [approx[2][0][0], approx[2][0][1]],
                          [approx[3][0][0], approx[3][0][1]]], dtype=np.float64)
    retval, rvec, tvec = cv2.solvePnP(objPoints, imgPoints, cameraMatrix, distCoeffs)
    R = cv2.Rodrigues(rvec)
    print("四个角点的坐标为：", approx)
    print("旋转向量为：", rvec)
    print("平移向量为：", tvec)
    print("旋转矩阵为：", R)


if __name__ == '__main__':
    cap = cv2.VideoCapture(0)
    while cap.isOpened():
        ret, frame = cap.read()

        approx = color(frame)

        c = cv2.waitKey(1)
        if c == ord('p'):
            pnp(approx)
        elif c == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

